import json
import tkinter as tk
from tkinter import filedialog, messagebox

import matplotlib.pyplot as plt
import numpy as np
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image


# 加载类别标签
def load_class_labels(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        return json.load(f)


# 加载模型
model = load_model('model/mushroom_classification_model_20241222_233510.keras')  # 修改为您的模型文件名
class_labels = load_class_labels('model/class_labels.json')  # 确保该文件存在


def predict_image(img_path, top_n=3):
    # 加载图像并进行预处理
    img = image.load_img(img_path, target_size=(150, 150))  # 根据您的模型输入大小调整
    img_array = image.img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)  # 归一化处理

    # 进行预测
    predictions = model.predict(img_array)

    # 获取前 top_n 个预测结果及其置信度
    top_indices = np.argsort(predictions[0])[-top_n:][::-1]
    top_classes = [(class_labels[i], predictions[0][i]) for i in top_indices]

    return top_classes


def open_file():
    file_path = filedialog.askopenfilename(
        title="选择图像文件",
        filetypes=[("Image files", "*.jpg;*.jpeg;*.png;*.bmp")]
    )
    if file_path:
        try:
            top_classes = predict_image(file_path, top_n=3)  # 获取前3个预测结果
            result_text = "\n".join([f"{cls}: {conf:.2f}" for cls, conf in top_classes])
            messagebox.showinfo("预测结果", f"预测结果:\n{result_text}")

            # 显示图像和预测结果
            plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体
            plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
            plt.imshow(image.load_img(file_path))
            plt.title(f"预测结果:\n{result_text}")
            plt.axis('off')  # 不显示坐标轴
            plt.show()
        except Exception as e:
            messagebox.showerror("错误", str(e))


# 创建 GUI 窗口
root = tk.Tk()
root.title("图像分类预测")

# 创建按钮
btn_open = tk.Button(root, text="选择图像文件", command=open_file)
btn_open.pack(pady=20)

# 运行 GUI 主循环
root.mainloop()
